#include <stdio.h>
#include "blas.h"
#include "softmax_layer.h"
#include "activations.h"
#include "network.h"


void forward_softmax_layer(const softmax_layer l, network net)
{
    // if(l.softmax_tree){
    //     int i;
    //     int count = 0;
    //     for (i = 0; i < l.softmax_tree->groups; ++i) {
    //         int group_size = l.softmax_tree->group_size[i];
    //         softmax_cpu(net.input + count, group_size, l.batch, l.inputs, 1, 0, 1, l.temperature, l.output + count);
    //         count += group_size;
    //     }
    // } else {
        softmax_cpu(net.input, l.inputs/l.groups, l.batch, l.inputs, l.groups, l.inputs/l.groups, 1, l.temperature, l.output);
    // }

    if (net.truth && !l.noloss) {
        softmax_x_ent_cpu(l.batch*l.inputs, l.output, net.truth, l.delta, l.loss);
        l.cost[0] = sum_array(l.loss, l.batch*l.inputs);
    }
}

void backward_softmax_layer(const softmax_layer l, network net)
{
    axpy_cpu(l.inputs*l.batch, 1, l.delta, 1, net.delta, 1);
}

softmax_layer make_softmax_layer(int batch, int inputs, int groups)
{
    softmax_layer l = {0};
    l.type = SOFTMAX;
    l.batch = batch;
    l.groups = groups;
    l.inputs = inputs;
    l.outputs = inputs;
    l.output = calloc(inputs * batch, sizeof(float));
    l.delta = calloc(inputs * batch, sizeof(float));
    l.loss = calloc(inputs * batch, sizeof(float));
    l.cost = calloc(1, sizeof(float));
    l.forward = forward_softmax_layer;
    l.backward = backward_softmax_layer;
    l.temperature = 1.0;

    return l;
}